"""
Copyright 2025 Google. This software is provided as-is, 
without warranty or representation for any use or purpose. 
Your use of it is subject to your agreement with Google.

"""

from google.cloud import bigquery
from google.cloud import storage
import google.api_core.exceptions
import config_vars


def export_bigquery_table_to_bq_gcs_staging(
    bq_client_main,
    input_bq_table_names,  # Pass the string of table names
    staging_path_bigquery,
    input_gcs_path,  # Pass the batch-specific path
):
    """
    Exports data from a BigQuery table to a GCS bucket in CSV format.

    """
    if config_vars.SOURCE_TYPE == "BigQuery":
        num_rows = 1000
        for table_name in input_bq_table_names.split(","):
            table_name = table_name.strip()  # Remove leading/trailing spaces

            # Construct the extract job configuration
            job_config = bigquery.ExtractJobConfig(
                destination_format="CSV",
                print_header=True,  # Include header row in the CSV
            )

            # Construct the destination uri
            destination_uri = f"{staging_path_bigquery}/{table_name.split('.')[-1]}.csv"  # Use table name for file

            # Create and run the extract job
            extract_job = bq_client_main.extract_table(
                source=table_name,
                destination_uris=destination_uri,
                job_config=job_config,
            )  # API request
            extract_job.result()  # Waits for job to complete

            print(
                f"Exported {num_rows if num_rows else 'all'} rows from {table_name} to {destination_uri}"
            )

        # Split the string into individual table names
        table_names = [name.strip() for name in input_bq_table_names.split(",")]

        for table_name in table_names:
            # Extract the table name without the project and dataset
            simple_table_name = table_name.split(".")[-1]
            input_gcs_path[simple_table_name] = (
                f"{staging_path_bigquery}/{simple_table_name}.csv"
            )

        # Print the resulting dictionary (optional)
        return input_gcs_path


def create_table_if_not_exists(source_table_id, target_table_id):
    """
    Checks if a target BigQuery table exists and creates it if it doesn't,
    using the schema of a source table.

    Args:
        source_table_id: The full ID of the source BigQuery table
                         in the format 'project_id.dataset_id.table_id'.
        target_table_id: The full ID of the target BigQuery table
                         in the format 'project_id.dataset_id.table_id'.
    """

    # Initialize BigQuery client
    client = bigquery.Client()

    # Check if the target table exists
    try:
        client.get_table(target_table_id)
    except google.api_core.exceptions.NotFound:
        print(f"Table {target_table_id} not found, creating it...")

        # Get the source table schema
        source_table = client.get_table(source_table_id)
        schema = source_table.schema

        # Create the target table with the source table's schema
        table = bigquery.Table(target_table_id, schema=schema)
        table = client.create_table(table)
        print(f"Table {target_table_id} created successfully.")


def import_gcs_to_bigquery(table_attributes):
    """
    Loads data from GCS to BigQuery. Performs INSERT operations only.
    Creates target table if it doesn't exist.

    Args:
        None

    Returns:
        None
    """

    client = bigquery.Client(project=config_vars.target_bq_project_id)
    storage.Client()

    for table_name in config_vars.user_counts:
        # Get GCS URI and table attributes
        gcs_uri = table_attributes[table_name]["output_gcs_path"]

        # Construct source table ID
        source_table_id = (
            config_vars.source_bq_project_id
            + "."
            + config_vars.source_bq_dataset
            + "."
            + table_name
        )

        # Construct target table ID
        target_table_id = f"{config_vars.target_bq_project_id}.{config_vars.target_bq_dataset}.{table_name}"

        # Create table if it doesn't exist
        create_table_if_not_exists(source_table_id, target_table_id)

        # Construct load job configuration
        job_config = bigquery.LoadJobConfig(
            source_format=bigquery.SourceFormat.CSV,
            skip_leading_rows=1,  # Skip header row if present
            autodetect=False,
            ignore_unknown_values=True,  # Ignore the id column generated by snowfakery
            write_disposition=bigquery.WriteDisposition.WRITE_APPEND,  # Append to existing data
        )

        # Load data from GCS to BigQuery
        load_job = client.load_table_from_uri(
            gcs_uri, target_table_id, job_config=job_config
        )
        load_job.result()  # Wait for job to complete

        print(f"Loaded data from {gcs_uri} to {target_table_id}")
